{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import reasoners.benchmark.bw_utils as utils\n",
    "import sys\n",
    "import os\n",
    "os.chdir('../..')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "config_file: str = \"examples/rap_blocksworld/data/bw_config.yaml\"\n",
    "domain_file: str = \"examples/rap_blocksworld/data/generated_domain.pddl\"\n",
    "prompt_path: str = 'examples/rap_blocksworld/prompts/prompt.json'\n",
    "data_path: str = 'examples/rap_blocksworld/data/new_data/step_4.json'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "47\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "86\n",
      "155\n",
      "153\n",
      "113\n",
      "46\n"
     ]
    }
   ],
   "source": [
    "for i in [2,4,6,8,10,12]:\n",
    "    data_path: str = f'examples/rap_blocksworld/data/full_data/step_{i}.json'\n",
    "    data_list = json.load(open(data_path, 'r'))\n",
    "    print(len(data_list))\n",
    "    random.shuffle(data_list)\n",
    "    prompt_set = utils.load_blocksworld(config_file, domain_file, data_list=data_list[:10], return_intermediate=True)\n",
    "    # dataset = utils.load_blocksworld(config_file, domain_file, data_list=data_list[10:], return_intermediate=False)\n",
    "    pool = [\n",
    "    {\"init\": p[\"init\"], \"goal\": p[\"goal\"], \"plan\": p[\"plan\"], \"states\": p[\"states\"]}\n",
    "    for p in prompt_set]\n",
    "    prompt_dict = json.load(open(\"examples/rap_blocksworld/prompts/pool_prompt.json\"))\n",
    "    prompt_dict[\"example_pool\"] = pool\n",
    "    json.dump(prompt_dict, open(f'examples/rap_blocksworld/prompts/pool_prompt_v2_step_{i}.json', 'w'))\n",
    "    json.dump(data_list[10:], open(f'examples/rap_blocksworld/data/split_v2_step_{i}_data.json', 'w'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "602\n"
     ]
    }
   ],
   "source": [
    "data_path: str = f'examples/rap_blocksworld/data/full.json'\n",
    "data_list = json.load(open(data_path, 'r'))\n",
    "print(len(data_list))\n",
    "random.shuffle(data_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_set = utils.load_blocksworld(config_file, domain_file, data_list=data_list[:10], return_intermediate=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "pool = [\n",
    "    {\"init\": p[\"init\"], \"goal\": p[\"goal\"], \"plan\": p[\"plan\"], \"states\": p[\"states\"]}\n",
    "    for p in prompt_set]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_dict = json.load(open(\"examples/rap_blocksworld/prompts/pool_prompt.json\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_dict[\"example_pool\"] = pool\n",
    "json.dump(prompt_dict, open('examples/rap_blocksworld/prompts/pool_prompt_v1.json', 'w'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[7, 11, 5, 5, 11, 9, 13, 9, 9, 7]"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[len(ex[\"plan\"].split(\"\\n\")) for ex in prompt_dict[\"example_pool\"]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['LLMs-Planning/llm_planning_analysis/instances/blocksworld/generated_basic_3/instance-34.pddl',\n",
       " '(unstack c b)\\n(put-down c)\\n(pick-up b)\\n(stack b c)\\n',\n",
       " 4]"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_list[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in [2, 4, 6, 8, 10, 12]:\n",
    "    step_data = [d for d in data_list[10:] if d[-1] == i]\n",
    "    json.dump(step_data, open(f'examples/rap_blocksworld/data/split_v1/split_v1_step_{i}_data.json', 'w'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "default",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.6"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
